#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
作  者    : 北极星光 light22@126.com
创建时间  : 2023-07-28 11:38:03
最后编辑  : 2024-03-12 23:13:54
编辑人    : 北极星光
'''

import requests
import yaml
from openai import OpenAI

class DeepseekAI:
    def __init__(self,api_key=None, model_name=None, bot_name=None, bot_role=None, welcome_words=None, sensitive_words=None):
        self.api_key = api_key
        self.model_name = model_name
        self.bot_name = bot_name
        self.bot_role = bot_role
        self.welcome_words = welcome_words
        self.sensitive_words = sensitive_words
        self.messages = {'default_user': [{'role':'system', 'content': self.bot_role}]}
        self.client = OpenAI(api_key=self.api_key, base_url="https://api.deepseek.com")

    # 对话
    def deepseek_chat(self, text,user='default_user'):
        if user not in self.messages:
            self.messages[user] = [{'role': 'system', 'content': self.bot_role}]
        self.messages[user].append({'role': 'user', 'content': text})
        # 过滤敏感词
        if any(word in text for word in self.sensitive_words):
            response = '''### Deepseek AI对话
            消息中含有敏感词，请重新输入。'''
            return response

        response = self.client.chat.completions.create(
            model=self.model_name,
            messages=self.messages[user]
        )
        self.messages[user].append({'role':response.choices[0].message.role, 'content': response.choices[0].message.content})
        reply_text = f'''### Deepseek AI对话
        {response.choices[0].message.content}'''

        return reply_text

    # 列出所有模型
    def list_models(self):
        response = self.client.models.list()
        model_list = [i.get('id') for i in response.to_dict().get('data')]
        model_list_text = f'''### Deepseek 所有模型列表'''
        for model in model_list:
            model_list_text += f'''
- {model}'''

        return model_list_text

    # 查询余额
    def get_balance(self):
        url = "https://api.deepseek.com/user/balance"

        payload={}
        headers = {
        'Accept': 'application/json',
        'Authorization': f'Bearer {self.api_key}'
        }

        response = requests.request("GET", url, headers=headers, data=payload)
        currency = response.json().get('balance_infos')[0].get('currency')
        total_balance = response.json().get('balance_infos')[0].get('total_balance')
        granted_balance = response.json().get('balance_infos')[0].get('granted_balance')
        topped_up_balance = response.json().get('balance_infos')[0].get('topped_up_balance')
        balance_text = f'''### 余额信息
- 货币单位：{currency}
- 总余额：{total_balance}
- 赠送余额：{granted_balance}
- 充值余额：{topped_up_balance}'''

        return balance_text
    
    # 重置消息记录
    def reset_messages(self, user='default_user'):
        self.messages[user] = [{'role': 'system', 'content': self.bot_role}]


if __name__ == '__main__':
    with open('./config/ds_config.yaml', 'r', encoding='utf-8') as f:
        cfg = yaml.load(f, Loader=yaml.FullLoader)
    text = "遐蝶的遗器选用散兵节奏还是狗屁不通生命球好呢还是量子球好呢？"
    deepseek_ai = DeepseekAI(api_key=cfg.get('api_key'), model_name=cfg.get('model_name'), bot_name=cfg.get('bot_name'), bot_role=cfg.get('bot_role'), welcome_words=cfg.get('welcome_words'), sensitive_words=cfg.get('sensitive_words'))
    response = deepseek_ai.deepseek_chat(text)
    print(response)
    print(deepseek_ai.messages)
    # print(deepseek_ai.list_models())
    # print(deepseek_ai.get_balance())